// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of SNP Sampling Functions for LD Calculations
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "copyfile.h"
#include "ld_data.h"
#include "suffixes.h"
#include "types.h"

// Function assumes that the map data has been sorted.
int ld_data_get_indices(struct small_locus *map, int nMap, char *snp_id, int disp_r_window_kb, int *start_index, int *end_index, int *anchor_index) {
	int i, snp_index = -1, window_bp = 100000, lower_bp_limit = 0, upper_bp_limit = 0;

	if (map == NULL || nMap <= 0 || snp_id == NULL || disp_r_window_kb <= 0 || start_index == NULL || end_index == NULL || anchor_index == NULL) return 0;
	else if (strlen(snp_id) == 0) return 0;

	window_bp = 1000 * disp_r_window_kb;

	for (i = 0; i < nMap; i++) {
		if (strcmp(map[i].name, snp_id) == 0) {
			snp_index = i;
			break;
		}
	}

	if (snp_index == -1) return 0;
	else if (map[snp_index].bp == 0) return 0; // Incase of EST which don't have a BP in a data-set.

	*anchor_index = snp_index;
	lower_bp_limit =  map[snp_index].bp - window_bp;
	upper_bp_limit =  map[snp_index].bp + window_bp;
	if (lower_bp_limit <= 0) lower_bp_limit = 1;

	*start_index = *end_index = -1;

	// Find the start index.
	for (i = snp_index; i >= 0; i--) {
		if (map[snp_index].chr != map[i].chr) break;
		else if (map[i].bp < lower_bp_limit) break;

		*start_index = i;
	}

	// Find the end index.
	for (i = snp_index; i < nMap; i++) {
		if (map[snp_index].chr != map[i].chr) break;
		else if (map[i].bp >= upper_bp_limit) break;

		*end_index = i + 1;
	}
	
	if (*start_index >= 0 && *end_index >= 0 && *start_index <= *end_index) return 1;
	else return 0;
}

int ld_data_create_tpedfile(struct selected_options *ops, char *bpedData, struct small_locus *map, int nMap, int start_snp_index, int end_snp_index, int nSamples) {
	int i, j;
	size_t snp_offset = 0, person_offset = 0, nSNPsWrote = 0;
	char *pBpedData = NULL, localFilepath[FILE_PATH_LENGTH];
	FILE *out = NULL;

	if (ops == NULL || bpedData == NULL || map == NULL || nMap <= 0 || start_snp_index <= 0 || end_snp_index <= 0 || nSamples <= 0) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0 || start_snp_index > end_snp_index) return 0;

	person_offset = 2 * sizeof(char);
	snp_offset = 2 * nSamples * sizeof(char);

	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_TPED_FILE);
	out = fopen(localFilepath, "w");
	if (out == NULL) return 0;
	for (i = start_snp_index, pBpedData = bpedData; i < end_snp_index; i++, bpedData += snp_offset) {
		int nFF = 0, nTT = 0;
		BOOL moreTT = FALSE;
		if (map[i].bp == 0) continue;

		for (j = 0, pBpedData = bpedData; j < nSamples; j++, pBpedData += person_offset) {
			if (pBpedData[0] == FALSE && pBpedData[1] == FALSE) nFF++;
			else if (pBpedData[0] == TRUE && pBpedData[1] == TRUE) nTT++;
		}

		if (nTT > nFF) moreTT = TRUE;

		fprintf(out, "%i %s 0 %i", map[i].chr, map[i].name, map[i].bp);
		if (moreTT) {
			for (j = 0, pBpedData = bpedData; j < nSamples; j++, pBpedData += person_offset) {
				// Reverse of the PLINK input behavior.
				if (pBpedData[0] == FALSE && pBpedData[1] == FALSE) fprintf(out, " 1 1");
				else if (pBpedData[0] == FALSE && pBpedData[1] == TRUE) fprintf(out, " 1 2");
				else if (pBpedData[0] == TRUE && pBpedData[1] == TRUE) fprintf(out, " 2 2");
				else if (pBpedData[0] == TRUE && pBpedData[1] == FALSE) fprintf(out, " %s %s", ops->missing_genotype, ops->missing_genotype);
			}
		} else {
			for (j = 0, pBpedData = bpedData; j < nSamples; j++, pBpedData += person_offset) {
				// Reverse of the PLINK input behavior.
				if (pBpedData[0] == FALSE && pBpedData[1] == FALSE) fprintf(out, " 2 2");
				else if (pBpedData[0] == FALSE && pBpedData[1] == TRUE) fprintf(out, " 1 2");
				else if (pBpedData[0] == TRUE && pBpedData[1] == TRUE) fprintf(out, " 1 1");
				else if (pBpedData[0] == TRUE && pBpedData[1] == FALSE) fprintf(out, " %s %s", ops->missing_genotype, ops->missing_genotype);
			}
		}
		fprintf(out, "\n");
		nSNPsWrote++;
	}
	fclose(out);

	if (nSNPsWrote == 0) {
		deleteFile(localFilepath);
		return 0;
	}

	return 1;
}

int ld_data_create_tfamfile(struct selected_options *ops, BOOL *sex, int nSex) {
	char localFilepath[FILE_PATH_LENGTH];
	FILE *out = NULL;
	int i;

	if (ops == NULL || sex == NULL || nSex <= 0) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0) return 0;

	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_TFAM_FILE);
	out = fopen(localFilepath, "w");
	if (out == NULL) return 0;
	for (i = 0; i < nSex; i++) {
		int gender_code = 2; // Female by default, not affected

		if (sex[i]) gender_code = 1;
		fprintf(out, "1 %i 0 0 %i -9\n", i + 1, gender_code);
	}
	fclose(out);

	return 1;
}

int ld_data_create_gphefile(struct selected_options *ops, BOOL *sex, int nSex) {
	char localFilepath[FILE_PATH_LENGTH];
	FILE *out = NULL;
	int i;

	if (ops == NULL || sex == NULL || nSex <= 0) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0) return 0;

	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_GPHE_FILE);
	out = fopen(localFilepath, "w");
	if (out == NULL) return 0;
	fprintf(out, "id sex phe\n");
	for (i = 0; i < nSex; i++) {
		int gender_code = 0; // Female by default, not affected

		if (sex[i]) gender_code = 1;
		fprintf(out, "%i %i NA\n", i + 1, gender_code);
	}
	fclose(out);

	return 1;
}

int ld_data_create_smapfile(struct selected_options *ops, struct small_locus *map, int nMap, int start_snp_index, int end_snp_index, int anchor_snp_index) {
	char localFilepath[FILE_PATH_LENGTH];
	FILE *out = NULL;
	int i;

	if (ops == NULL || map == NULL || nMap <= 0 || start_snp_index <= 0 || end_snp_index <= 0) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0 || start_snp_index > end_snp_index) return 0;

	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_SMAP_FILE);
	out = fopen(localFilepath, "w");
	if (out == NULL) return 0;

	fprintf(out, "SNP BP\n");
	for (i = start_snp_index; i < end_snp_index; i++) {
		if (i != anchor_snp_index) fprintf(out, "%s %i\n", map[i].name, map[i].bp);
		else fprintf(out, "<<<<<<<%s %i\n", map[i].name, map[i].bp);
	}

	return 1;
}

int ld_data_get_indices_nmer(struct small_locus *map, int nMap, char *snp_id, int window_element_size, int *start_index, int *end_index, int *anchor_index) {
	int i, j, snp_index = -1;

	if (map == NULL || nMap <= 0 || snp_id == NULL || window_element_size <= 0 || start_index == NULL || end_index == NULL) return 0;
	else if (strlen(snp_id) == 0) return 0;

	for (i = 0; i < nMap; i++) {
		if (strcmp(map[i].name, snp_id) == 0) {
			snp_index = i;
			break;
		}
	}

	if (snp_index == -1) return 0;
	else if (map[snp_index].bp == 0) return 0; // Incase of EST which don't have a BP in a data-set.

	*anchor_index = snp_index;
	*start_index = *end_index = -1;
	for (i = *anchor_index - 1, j = 0; i >= 0 && j < window_element_size; i--, j++) {
		if (map[snp_index].chr != map[i].chr) break;
		*start_index = i;
	}
	for (i = *anchor_index + 1, j = 0; i < nMap && j <= window_element_size; i++, j++) {
		if (map[snp_index].chr != map[i].chr) break;
		*end_index = i;
	}
	
	if (*start_index >= 0 && *end_index >= 0 && *start_index <= *end_index) return 1;
	else return 0;
}
